import argparse
import functools
import subprocess
import hashlib
import radix
from decimal import Decimal
import os
from ipaddress import IPv4Network
import random

from SimulationConfiguration import Conf
from process_log_file import get_incoming_packets
from analyze_ssfnet_output import analyze_attack
from TrafficGenerator.traceroute_utils import unix_process_tracert_path

ttl_penalty = {15169: 3, 2603: 1, 31815: 2, 1680: 4, 4323 : 1, 558 : 1, 14742: 2, 8283: 1, 3356 : 2, 15133: 1, 32475: 3, 10912: 3, 7843: 1, 11172: 1, 20965: 2, 22822: 3, 3561: 2, 26347: 2, 174: 7, 12849: 1,19271:2,1299:3, 378: 7, 701: 2, 36351: 4}

class PacketHeader:
    def __init__(self,flow_attrs, time, sip, sp, dip, dp, ttl):
        self.flow_attrs = flow_attrs
        self.time = float(time)
        self.sip = str(sip)
        self.sp = sp
        self.dip = str(dip)
        self.dp = dp
        self.ttl = abs(int(ttl))
    def __str__(self):
        return "time: {0} source = ({1}, {2}) dest = ({3},{4}) ttl = {5}".format(self.time, self.sip, self.sp, self.dip, self.dp, self.ttl)
    def __repr__(self):
        return self.__str__()
    def __hash__(self):
        m = hashlib.md5()
        for attr in self.flow_attrs:
            attr_value = self.__getattribute__(attr)
            m.update(str(attr_value).encode('utf-8'))
        return int(m.hexdigest(), 16)
    def __eq__(self, other):
        for attr in self.flow_attrs:
            left_value = self.__getattribute__(attr)
            right_value = other.__getattribute__(attr)
            if (left_value != right_value):
                return False
        return True
        
class IPAlgorithm:
    def __init__(self, flow_attrs, window):
            self.traffic = []
            self.flow_attrs = flow_attrs
            self.window = window
    def add_packet(self,time, ttl, sip, sp, dip, dp):
        self.traffic.append(PacketHeader(self.flow_attrs, time, sip, sp, dip, dp, ttl))
    def flows(self, starting_from_time, ending_at_time):
        flows = {}
        for packet in filter(lambda x: x.time >=starting_from_time and x.time <= ending_at_time, self.traffic):
            if packet in flows:
                flows[packet] = (flows[packet][0] + 1, flows[packet][1] + [packet])
            else:
                flows[packet] = (0, [packet])
        
        return flows
    
    def count_flows(self, starting_from_time = Conf.TRAFFIC_START_TIME, ending_at_time = Decimal('Infinity')):
        return len(self.flows(starting_from_time, ending_at_time))
        
    def get_flows_TTL_list(self, starting_from_time = Conf.TRAFFIC_START_TIME, ending_at_time = Decimal('Infinity')):
        return functools.reduce(lambda x,y: x + [y.ttl], self.flows(starting_from_time, ending_at_time), [])
        
    def count_TTL_flows_sum(self, starting_from_time = Conf.TRAFFIC_START_TIME, ending_at_time = Decimal('Infinity')):
        return sum(self.get_flows_TTL_list(starting_from_time, ending_at_time))
    
    def max_traffic_time(self):
        return max(map(lambda p: p.time, self.traffic))
    
    def min_traffic_time(self):
        return min(map(lambda p: p.time, self.traffic))
    
    DUMP_PARSER = 'tcpdump.jar'
    DUMP_PARSER_ARGS = 'time ttl sip sp dip dp'

INCOMING_INTERFACE = 'Lan3'
INTERFACE_COLUMN = 3

prefix_radix = radix.Radix()


def prepare_ip_mapping(prefix_to_asn_file, relevant_ases):
    with open(prefix_to_asn_file,'r') as origin_file:
        lines = origin_file.readlines()
    for line in lines:
        lineSplit = line.split()
        asn = int(lineSplit[1])
        if asn in relevant_ases:
            rnode = prefix_radix.add(lineSplit[0])
            rnode.data["asn"] = asn


def getASNofIP(ip):
    rnode = prefix_radix.search_best(ip)
    if rnode is not None:
        return rnode.data['asn']

    return None


def generatre_ttl_penalty():
    # Randomize a border TTL penalty
    border_ttl_penalty_distibution, border_ttl_penalty_distibution_params = \
        (Conf.TIER_1_BORDER_TTL_PENALTY_DISTRIBUTION, Conf.TIER_1_BORDER_TTL_PENALTY_DISTRIBUTION_ARGUMENTS)
    border_ttl_penalty = max(1, round(border_ttl_penalty_distibution(*border_ttl_penalty_distibution_params)))
    return border_ttl_penalty


def analyze_firewall_files(firewall_log_file, window, SSFNET_output_file, prefix_to_asn_file, attacker, victim, **kwargs):
    algorithm = IPAlgorithm(Conf.FLOW_ATTRIBUTES, window)

    attack = None
    attack_succeeded = False

    if SSFNET_output_file is not None:
        attack = analyze_attack(SSFNET_output_file, None, attacker, victim)
        attack_succeeded = not attacker in attack
        prepare_ip_mapping(prefix_to_asn_file, attack)

    packets, _= get_incoming_packets('firewall', firewall_log_file)


    for packet in packets:
        ttl_diff = 0
        if attack is not None and attack_succeeded:
            try:
                asn = getASNofIP(packet.src)
                if asn in attack:
                    index = attack.index(asn)
                    if attack[index].was_affected_by_attack():
                        # Assigning a different TTL penalty for each AS hop
                        ttl_diff = sum(map(lambda _: generatre_ttl_penalty(), range(attack[index].get_AS_hop_delta())))
            except Exception:
                # packet.src might not be an IP
                pass

        algorithm.add_packet(packet.time.timestamp(), packet.ttl + ttl_diff, packet.src, packet.src_port, packet.dst, packet.dst_port)
    
    main(algorithm, window, **kwargs)
    
def analyze_dump_files(dump_files, window, **kwargs):
    algorithm = IPAlgorithm(Conf.FLOW_ATTRIBUTES, window)
    for file in dump_files:
        p = subprocess.Popen('java -jar {0} {1} {2}'.format(algorithm.DUMP_PARSER, file, algorithm.DUMP_PARSER_ARGS), stdout = subprocess.PIPE)
        print (file)
        for line in p.stdout.readlines():
            packet_data = line.decode("utf-8").split()
            algorithm.add_packet(*packet_data)

    main(algorithm, window, **kwargs)


def simulate_dump_files(dump_file, t0, t1, window, SSFNET_output_file, prefix_to_asn_file, attacker, victim, **kwargs):
    attack = None
    attack_succeeded = False

    if SSFNET_output_file is not None:
        attack = analyze_attack(SSFNET_output_file, None, attacker, victim)
        attack_succeeded = not attacker in attack
        relevant_ases_for_ip_mapping = attack.copy() + [int(attacker),int(victim)]
        prepare_ip_mapping(prefix_to_asn_file, relevant_ases_for_ip_mapping)

    packets, _= get_incoming_packets('tcpdump', dump_file)

    algorithm = IPAlgorithm(Conf.FLOW_ATTRIBUTES, window)

    # Find an IP of the attacker
    for node in prefix_radix.nodes():
        if node.data['asn'] == int(attacker):
            network = IPv4Network(node.prefix)
            address = random.choice(list(network.hosts()))
            break
    ttl_penalty_from_attacker_to_victim = len(unix_process_tracert_path(address.exploded))
    ttl_penalty_average = sum(ttl_penalty.values())/len(ttl_penalty)
    tracert_cache = {}
    for packet in packets:

        ttl = packet.ttl
        if round(packet.time) in range(t0, t1):
            if attack is not None and attack_succeeded:
                try:
                    asn = getASNofIP(packet.src.exploded)
                    if asn in attack:
                        index = attack.index(asn)
                        if attack[index].was_affected_by_attack():
                            # Assigning a different TTL penalty for each AS hop
                            # Traceroute to destination to determine the initial TTL value
                            if not packet.src.exploded in tracert_cache:
                                tracert_len_to_destination = len(unix_process_tracert_path(packet.src.exploded))
                                tracert_cache[packet.src.exploded] = tracert_len_to_destination
                            else:
                                print('Tracert cache hit for %s' % (packet.src.exploded,))
                                tracert_len_to_destination = tracert_cache[packet.src.exploded]
                            initial_ttl_value = ttl + tracert_len_to_destination
                            new_ttl_penalty = sum(map(lambda asn: ttl_penalty[int(asn)] if int(asn) in ttl_penalty else ttl_penalty_average, attack[index].get_path_to_victim())) + \
                                    ttl_penalty_from_attacker_to_victim
                            ttl = initial_ttl_value - int(new_ttl_penalty)
                except Exception as e:
                    # packet.src might not be an IP
                    print("Error getting asn of %s: %s" % (packet.src, e))
                    pass

        algorithm.add_packet(packet.time, ttl, packet.src.exploded, packet.src_port, packet.dst.exploded, packet.dst_port)

    main(algorithm, window, **kwargs)
            
def main(ip_algorithm, window, **kwargs):
    flow_counts = []
    ttl_sums = []
    ttl_averages = []
    max_traffic_time = int(ip_algorithm.max_traffic_time())
    min_traffic_time = int(ip_algorithm.min_traffic_time())
    if not window:
        window = max_traffic_time - min_traffic_time
    for time in range(min_traffic_time, max_traffic_time, window):
        flow_counts.append(ip_algorithm.count_flows(time, time + window))
        ttl_sums.append(ip_algorithm.count_TTL_flows_sum(time, time + window))
        ttl_averages.append(round(ttl_sums[-1] / flow_counts[-1], 2))
    print("Flows over Windows:\n{0}\n".format(flow_counts))
    print("TTL sums over Windows:\n{0}\n".format(ttl_sums))
    print("TTL averages over Windows:\n{0}\n".format(ttl_averages))
    with open('results.txt','w') as results:
        for x in flow_counts[:10]:
            results.write(str(x) + '\t')
        results.write(os.linesep)
        for x in flow_counts[10:]:
            results.write(str(x) + '\t')        
        results.write(os.linesep)
        for x in ttl_sums[:10]:
            results.write(str(x) + '\t')        
        results.write(os.linesep)
        for x in ttl_sums[10:]:
            results.write(str(x) + '\t')        
        results.write(os.linesep)
        for x in ttl_averages[:10]:
            results.write(str(x) + '\t')        
        results.write(os.linesep)
        for x in ttl_averages[10:]:
            results.write(str(x) + '\t')        
        results.write(os.linesep)
    

if (__name__ == "__main__"):
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers()
    
    # TCPDump files subparser
    tcpdump = subparsers.add_parser('tcpdump', help='Parses TCPDump files')
    tcpdump.set_defaults(func=analyze_dump_files)
    tcpdump.add_argument('dump_files', nargs='+', type=str, help='List of TCPDump files.')

    # TCPDump files to with attack simulation subparser
    tcpdump_simulate = subparsers.add_parser('simulate_tcpdump', help='Parses TCPDump files and simulate attack')
    tcpdump_simulate.set_defaults(func=simulate_dump_files)
    tcpdump_simulate.add_argument('dump_file', type=str, help='A TCPDump file.')
    tcpdump_simulate.add_argument('t0', type = int)
    tcpdump_simulate.add_argument('t1', type = int)
    tcpdump_simulate.add_argument("-s", "--SSFNET_output_file", help="SSFNET execution textual output file")
    tcpdump_simulate.add_argument("-p", "--prefix_to_asn_file", help="Mapping of IPv4 Prefixes to asns")
    tcpdump_simulate.add_argument("-a", "--attacker", help="The attacker AS")
    tcpdump_simulate.add_argument("-v", "--victim", help="The victim AS")
    
    # Firewall Log files subparser
    firewall = subparsers.add_parser('firewall', help='Parses firewall log files')
    firewall.set_defaults(func=analyze_firewall_files)
    firewall.add_argument('firewall_log_file', type=str, help='A firewall log file path')
    firewall.add_argument("-s", "--SSFNET_output_file", help="SSFNET execution textual output file")
    firewall.add_argument("-p", "--prefix_to_asn_file", help="Mapping of IPv4 Prefixes to asns")
    firewall.add_argument("-a", "--attacker", help="The attacker AS")
    firewall.add_argument("-v", "--victim", help="The victim AS")

    parser.add_argument('-w', '--window', type = int)
    args = parser.parse_args()
    
    args.func(**vars(args))